#!/usr/bin/env python3
import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import threading
import cv2
import numpy as np
import time
import os
import json
import queue
from collections import defaultdict
from ultralytics import YOLO
from datetime import datetime
import os
os.environ['LD_PRELOAD'] = '/usr/lib/aarch64-linux-gnu/libgomp.so.1'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OMP_STACKSIZE'] = '4M'
# 配置参数
SAVE_DIR = "./results"
SEG_MODEL_PATH = "./merge_seg_epoch201_rknn_model"
POSE_MODEL_PATH = "./merge_pose_epoch600_rknn_model"
SEG_CORE = 0  # 分割模型专用核心
POSE_CORE = 1  # 关键点模型专用核心

# 位置计算类（示例）
class ComputePosition:
    def __init__(self):
        # 初始化参数
        self.calibration_params = {}
        rospy.loginfo("位置计算器初始化完成")
    
    def get_position(self, color_img, depth_img, seg_results, pose_results):
        """
        计算最终位置信息
        :param color_img: 彩色图像
        :param depth_img: 深度图像
        :param seg_results: 分割结果
        :param pose_results: 关键点结果
        :return: 位置信息字典
        """
        # 这里实现具体的位置计算逻辑
        # 示例：返回一个包含位置信息的字典
        position_info = {
            "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S_%f"),
            "x": 0.0,
            "y": 0.0,
            "z": 0.0,
            "rotation": 0.0,
            "confidence": 0.95
        }
        return position_info

class NPUProcessor:
    def __init__(self, core_id, model_path, task_type):
        self.core_id = core_id
        self.task_type = task_type
        self.task_queue = queue.Queue(maxsize=20)  # 任务队列
        self.result_queue = None  # 结果队列（由主控制器设置）
        self.running = True
        
        # 加载模型
        self.load_model(model_path, task_type)
        
        # 启动处理线程
        self.thread = threading.Thread(target=self.process_tasks)
        self.thread.daemon = True
        self.thread.start()
    
    def load_model(self, model_path, task_type):
        """加载模型到指定核心"""
        # 设置NPU核心绑定
        os.environ['RKNN_SERVER_CORE_MASK'] = str(1 << self.core_id)
        
        # 加载模型
        self.model = YOLO(model_path, task=task_type)
        rospy.loginfo(f"核心 {self.core_id}: {task_type}模型加载完成")
    
    def add_task(self, color_img, depth_img, timestamp):
        """添加任务到队列"""
        self.task_queue.put({
            "color": color_img,
            "depth": depth_img,
            "timestamp": timestamp
        })
    
    def process_tasks(self):
        """处理任务队列中的任务"""
        while self.running or not self.task_queue.empty():
            try:
                # 获取任务
                task = self.task_queue.get(timeout=0.1)
                color_img = task["color"]
                timestamp = task["timestamp"]
                
                # 处理任务
                results = self.model.predict(color_img, save=False, verbose=False)
                
                # 检查结果
                success = False
                result_img = None
                
                if self.task_type == "segment":
                    if results and results[0].masks is not None:
                        result_img = results[0].plot()
                        success = True
                else:  # pose
                    if results and results[0].keypoints is not None:
                        result_img = results[0].plot()
                        success = True
                
                # 发送结果
                if self.result_queue:
                    self.result_queue.put({
                        "type": self.task_type,
                        "timestamp": timestamp,
                        "result_img": result_img,
                        "raw_results": results,  # 添加原始结果
                        "success": success
                    })
                
            except queue.Empty:
                time.sleep(0.01)
            except Exception as e:
                rospy.logerr(f"核心 {self.core_id} 处理错误: {e}")

    def stop(self):
        """停止处理器"""
        self.running = False
        self.thread.join()

class RealTimeProcessor:
    def __init__(self):
        self.bridge = CvBridge()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        self.image_queue = queue.Queue(maxsize=20)  # 输入图像队列
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)
        rospy.loginfo("订阅了相机彩色和深度图像话题")
        
        # 创建NPU处理器
        self.seg_processor = NPUProcessor(SEG_CORE, SEG_MODEL_PATH, "segment")
        self.pose_processor = NPUProcessor(POSE_CORE, POSE_MODEL_PATH, "pose")
        
        # 位置计算器
        self.position_calculator = ComputePosition()
        
        # 结果队列
        self.result_queue = queue.Queue()
        
        # 为处理器设置结果队列
        self.seg_processor.result_queue = self.result_queue
        self.pose_processor.result_queue = self.result_queue
        
        # 结果整合器
        self.result_integrator_thread = threading.Thread(target=self.integrate_results)
        self.result_integrator_thread.daemon = True
        self.result_integrator_thread.start()
        
        # 位置计算线程池
        self.position_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
        self.position_futures = []
        
        # 性能监控
        self.start_time = time.time()
        self.processed_count = 0
        self.total_images = 0
        
        # 结果存储
        self.pending_results = {}  # 等待配对的结果 {timestamp: {"seg": result, "pose": result}}
        self.completed_results = {}  # 已完成的结果
        
        # 创建保存目录
        os.makedirs(SAVE_DIR, exist_ok=True)
        self.seg_save_dir = os.path.join(SAVE_DIR, "seg_results")
        self.pose_save_dir = os.path.join(SAVE_DIR, "pose_results")
        self.position_save_dir = os.path.join(SAVE_DIR, "position_results")
        os.makedirs(self.seg_save_dir, exist_ok=True)
        os.makedirs(self.pose_save_dir, exist_ok=True)
        os.makedirs(self.position_save_dir, exist_ok=True)

    def color_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            with self.lock:
                self.color_img = cv_image
                # 添加到处理队列
                if self.depth_img is not None:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
                    self.image_queue.put((self.color_img.copy(), self.depth_img.copy(), timestamp))
                    self.total_images += 1
        except Exception as e:
            rospy.logerr(f"彩色图像转换错误: {e}")

    def depth_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
            with self.lock:
                self.depth_img = cv_image
                # 添加到处理队列
                if self.color_img is not None:
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
                    self.image_queue.put((self.color_img.copy(), self.depth_img.copy(), timestamp))
                    self.total_images += 1
        except Exception as e:
            rospy.logerr(f"深度图像转换错误: {e}")

    def distribute_tasks(self):
        """分发任务到NPU处理器"""
        while not rospy.is_shutdown():
            try:
                # 从队列获取图像
                color_img, depth_img, timestamp = self.image_queue.get(timeout=0.1)
                
                # 分配分割任务
                self.seg_processor.add_task(color_img, depth_img, timestamp)
                
                # 分配关键点任务
                self.pose_processor.add_task(color_img, depth_img, timestamp)
                
            except queue.Empty:
                time.sleep(0.01)
            except Exception as e:
                rospy.logerr(f"任务分发错误: {e}")

    def integrate_results(self):
        """整合分割和关键点结果"""
        while not rospy.is_shutdown():
            try:
                result = self.result_queue.get(timeout=0.1)
                timestamp = result["timestamp"]
                
                # 初始化结果记录
                if timestamp not in self.pending_results:
                    self.pending_results[timestamp] = {
                        "seg": None, 
                        "pose": None,
                        "color": None,
                        "depth": None
                    }
                
                # 保存结果
                if result["type"] == "segment":
                    self.pending_results[timestamp]["seg"] = result
                elif result["type"] == "pose":
                    self.pending_results[timestamp]["pose"] = result
                
                # 检查是否配对完成
                if (self.pending_results[timestamp]["seg"] is not None and 
                    self.pending_results[timestamp]["pose"] is not None):
                    
                    seg_result = self.pending_results[timestamp]["seg"]
                    pose_result = self.pending_results[timestamp]["pose"]
                    
                    # 保存结果图像
                    if seg_result["success"]:
                        cv2.imwrite(os.path.join(self.seg_save_dir, f"{timestamp}.png"), seg_result["result_img"])
                    
                    if pose_result["success"]:
                        cv2.imwrite(os.path.join(self.pose_save_dir, f"{timestamp}.png"), pose_result["result_img"])
                    
                    # 获取原始图像
                    color_img = None
                    depth_img = None
                    try:
                        # 查找原始图像（可能已从队列中移除）
                        # 在实际应用中，可能需要缓存原始图像
                        # 这里简化处理，使用结果中的时间戳作为标识
                        color_img = seg_result["color"]  # 假设任务中包含原始图像
                        depth_img = seg_result["depth"]
                    except:
                        rospy.logwarn(f"无法获取原始图像: {timestamp}")
                    
                    # 提交位置计算任务
                    if color_img is not None and depth_img is not None:
                        future = self.position_executor.submit(
                            self.calculate_position,
                            color_img,
                            depth_img,
                            seg_result["raw_results"],
                            pose_result["raw_results"],
                            timestamp
                        )
                        self.position_futures.append(future)
                    
                    # 移动到已完成
                    self.completed_results[timestamp] = {
                        "seg": seg_result["success"],
                        "pose": pose_result["success"]
                    }
                    
                    # 更新性能统计
                    self.processed_count += 1
                    elapsed = time.time() - self.start_time
                    fps = self.processed_count / elapsed if elapsed > 0 else float('inf')
                    
                    rospy.loginfo(f"处理完成: {timestamp} | "
                                 f"进度: {self.processed_count}/{self.total_images} | "
                                 f"FPS: {fps:.2f}")
                    
                    # 从待处理中移除
                    del self.pending_results[timestamp]
                
            except queue.Empty:
                time.sleep(0.01)
            except Exception as e:
                rospy.logerr(f"结果整合错误: {e}")

    def calculate_position(self, color_img, depth_img, seg_results, pose_results, timestamp):
        """计算位置信息并保存结果"""
        try:
            # 调用位置计算器
            position_info = self.position_calculator.get_position(
                color_img, depth_img, seg_results, pose_results
            )
            
            # 保存位置信息
            with open(os.path.join(self.position_save_dir, f"{timestamp}.json"), "w") as f:
                json.dump(position_info, f, indent=2)
            
            rospy.loginfo(f"位置计算完成: {timestamp} | 位置: {position_info}")
            return position_info
        except Exception as e:
            rospy.logerr(f"位置计算错误: {e}")
            return None

    def run(self):
        """主处理循环"""
        # 启动任务分发线程
        distributor_thread = threading.Thread(target=self.distribute_tasks)
        distributor_thread.daemon = True
        distributor_thread.start()
        
        # 主循环
        rate = rospy.Rate(10)  # 10Hz
        while not rospy.is_shutdown():
            # 实时显示队列状态
            seg_qsize = self.seg_processor.task_queue.qsize()
            pose_qsize = self.pose_processor.task_queue.qsize()
            rospy.loginfo_throttle(1, f"队列状态: 分割={seg_qsize}, 关键点={pose_qsize} | 待处理: {self.image_queue.qsize()}")
            
            # 清理已完成的位置计算任务
            self.position_futures = [f for f in self.position_futures if not f.done()]
            
            rate.sleep()

    def shutdown(self):
        """关闭时保存性能报告"""
        # 停止所有处理器
        self.seg_processor.stop()
        self.pose_processor.stop()
        
        # 关闭位置计算线程池
        self.position_executor.shutdown(wait=True)
        
        # 计算性能指标
        total_time = time.time() - self.start_time
        fps = self.processed_count / total_time if total_time > 0 else float('inf')
        
        # 保存性能报告
        report = {
            "total_images": self.total_images,
            "processed_images": self.processed_count,
            "total_time": total_time,
            "fps": fps,
            "npu_cores_used": [SEG_CORE, POSE_CORE]
        }
        
        with open(os.path.join(SAVE_DIR, "performance.json"), "w") as f:
            json.dump(report, f, indent=2)
        
        rospy.loginfo("\n" + "="*50)
        rospy.loginfo(f"处理完成: {self.processed_count}/{self.total_images}张图像")
        rospy.loginfo(f"总耗时: {total_time:.2f}秒")
        rospy.loginfo(f"平均帧率(FPS): {fps:.2f}")
        rospy.loginfo("="*50)

if __name__ == "__main__":
    rospy.init_node("real_time_yolo_processor", anonymous=True)
    processor = RealTimeProcessor()
    
    # 注册关闭钩子
    rospy.on_shutdown(processor.shutdown)
    
    try:
        processor.run()
    except rospy.ROSInterruptException:
        pass